package com.sunzm.flink.datastream.scala.state

import com.alibaba.fastjson.JSON
import org.apache.flink.api.common.state.{BroadcastState, MapStateDescriptor, ReadOnlyBroadcastState}
import org.apache.flink.streaming.api.datastream.BroadcastStream
import org.apache.flink.streaming.api.functions.co.KeyedBroadcastProcessFunction
import org.apache.flink.streaming.api.scala._
import org.apache.flink.util.Collector
import org.slf4j.{Logger, LoggerFactory}

/**
 *
 * ScalaStateDemo广播状态示例
 *
 * @author Administrator
 * @version 1.0
 * @date 2021-06-22 20:48
 */
object ScalaBroadcastStateDemo {
  private val logger: Logger = LoggerFactory.getLogger(this.getClass.getName.stripSuffix("$"))
  private val isLocal = true

  def main(args: Array[String]): Unit = {

    //1.创建执行的环境
    val env: StreamExecutionEnvironment = if (isLocal) {
      StreamExecutionEnvironment.createLocalEnvironmentWithWebUI()
    } else {
      StreamExecutionEnvironment.getExecutionEnvironment
    }

    //用户数据
    val userserDS: DataStream[User] = env.fromElements("zs,17,北京", "ls,18,上海", "ww,19,深圳")
      .map(line => {
        val fields = line.split(",")

        User(fields(0), fields(1).toInt, fields(2))
      })

    //将用户数据作为广播状态使用
    // 一个 map descriptor，它描述了用于用户名与用户本身的 map 存储结构
    val userStateDescriptor: MapStateDescriptor[String, User] = new MapStateDescriptor[String, User](
      "userBroadcastState",
      classOf[String],
      classOf[User])

    // 广播流，广播用户信息并且创建 broadcast state
    val userBroadcastStream: BroadcastStream[User] = userserDS.broadcast(userStateDescriptor)

    //订单数据流
    //{"orderId":"od1001", "userName":"zs"}
    val orderDS: DataStream[String] = env.socketTextStream("82.156.210.70", 9999)

    val keyedOrderDS: KeyedStream[Order, String] = orderDS.map(line => {
      val jSONObject = JSON.parseObject(line)

      val orderId = jSONObject.getString("orderId")
      val userName = jSONObject.getString("userName")

      Order(orderId, userName)
    }).keyBy(_.orderId)

    //将两个流进行连接
    val bdConnectedDS: BroadcastConnectedStream[Order, User] = keyedOrderDS.connect(userBroadcastStream)

    val connectedDS: DataStream[String] = bdConnectedDS.process(new KeyedBroadcastProcessFunction[String, Order, User, String] {
      /**
       * 处理非广播流中的每个数据
       *
       * @param value
       * @param ctx
       * @param out
       */
      override def processElement(value: Order,
                                  ctx: KeyedBroadcastProcessFunction[String, Order, User, String]#ReadOnlyContext,
                                  out: Collector[String]): Unit = {

        //获取广播状态
        val bdState: ReadOnlyBroadcastState[String, User] = ctx.getBroadcastState(userStateDescriptor)

        //判断是否包含这个用户
        val order: Order = value
        if (bdState.contains(order.userName)) {
          val user: User = bdState.get(order.userName)
          out.collect(s"关联成功: [订单ID]: ${order.orderId}, [用户名]: ${user.name}, [年龄]:${user.age}, [地址]: ${user.addr}")
        } else {
          out.collect(s"未关联到用户信息: [订单ID]: ${order.orderId}, [用户名]: ${order.userName} ")
        }
      }

      /**
       * 处理广播流中的每个数据
       *
       * @param value
       * @param ctx
       * @param out
       */
      override def processBroadcastElement(value: User,
                                           ctx: KeyedBroadcastProcessFunction[String, Order, User, String]#Context,
                                           out: Collector[String]): Unit = {
        val user: User = value

        logger.info(s"收到用户信息: ${user.age}, ${user.age}, ${user.addr}")

        //获取广播状态
        val bdState: BroadcastState[String, User] = ctx.getBroadcastState(userStateDescriptor)

        bdState.put(user.name, user)
      }
    })

    connectedDS.print()

    //5.执行
    env.execute(this.getClass.getSimpleName.stripSuffix("$"))
  }

  case class User(name: String, age: Int, addr: String)
  case class Order(orderId: String, userName: String)

}
